import numpy as np
import os
import os.path as osp
import math
from pathlib import Path
import imageio
from concurrent.futures import ProcessPoolExecutor
from base.ssim import ssim
from base.psnr import psnr


def find_template(path_config: dict, size_config: dict, sliding_step: int,
                  mode: str):
    assert mode in ['psnr', 'ssim']
    img_truth_path = path_config['img_truth_path']
    img_better_path = path_config['img_better_path']
    img_worse_path = path_config['img_worse_path']
    save_path = path_config['save_path']
    low_bound = -1

    sub_width = size_config['sub_width']
    sub_height = size_config['sub_height']

    img_truth = imageio.imread(img_truth_path)
    img_worse = imageio.imread(img_worse_path)
    img_better = imageio.imread(img_better_path)
    h, w, c = img_truth.shape

    if mode == 'psnr':
        criterion = psnr
    elif mode == 'ssim':
        criterion = ssim
    # 确定 low_bound
    if low_bound == -1:
        diff_list = []
        diff = 0
        for i in range(0, h - sub_height, sliding_step):
            for j in range(0, w - sub_width, sliding_step):
                sub_img_worse = img_worse[i:i + sub_height, j:j + sub_width, :]
                sub_img_better = img_better[i:i + sub_height,
                                            j:j + sub_width, :]
                sub_img_truth = img_truth[i:i + sub_height, j:j + sub_width, :]
                matrix_worse = criterion(sub_img_worse, sub_img_truth)
                matrix_better = criterion(sub_img_better, sub_img_truth)
                diff = matrix_better - matrix_worse
                if (diff > 0):
                    diff_list.append(diff)
        diff_len = len(diff_list)
        diff_list = sorted(diff_list, reverse=True)
        index = 10 if diff_len > 10 else diff_len
        if diff_len > 0:
            low_bound = diff_list[index]
        else:
            print("Not found.")
            return
    # 寻找目标区域
    max_diff = 0
    index = 0
    target_path = "Not found"
    for i in range(0, h - sub_height, sliding_step):
        for j in range(0, w - sub_width, sliding_step):
            sub_img_worse = img_worse[i:i + sub_height, j:j + sub_width, :]
            sub_img_better = img_better[i:i + sub_height, j:j + sub_width, :]
            sub_img_truth = img_truth[i:i + sub_height, j:j + sub_width, :]
            matrix_worse = criterion(sub_img_worse, sub_img_truth)
            matrix_better = criterion(sub_img_better, sub_img_truth)
            if matrix_better - matrix_worse > low_bound:
                target_path = img_truth_path
                target_truth = sub_img_truth
                target_worse = sub_img_worse
                target_better = sub_img_better

                common_save_path = str(save_path) + '/' + str(index)

                os.makedirs(common_save_path, exist_ok=True)
                imageio.imsave(
                    str(common_save_path) + "/truth-temp.png", target_truth)
                imageio.imsave(
                    str(common_save_path) + "/worse.png", target_worse)
                imageio.imsave(
                    str(common_save_path) + "/better.png", target_better)
                index += 1
                if matrix_better - matrix_worse > max_diff:
                    max_diff = matrix_better - matrix_worse
                    best_save_path = str(save_path) + '/best'
                    os.makedirs(best_save_path, exist_ok=True)
                    imageio.imsave(
                        str(best_save_path) + "/truth-temp.png", target_truth)
                    imageio.imsave(
                        str(best_save_path) + "/worse.png", target_worse)
                    imageio.imsave(
                        str(best_save_path) + "/better.png", target_better)
    print('{0:} {1:.6f} {2:}'.format(target_path, max_diff, index))


if __name__ == '__main__':
    # 参数设置
    mode = 'ssim'  # psnr 或 ssim
    sliding_step = 25
    sub_width = 100
    sub_height = 100
    dir_truth = 'find_template/before/HR/9'
    dir_worse = 'find_template/before/ovsr10k/9'
    dir_better = 'find_template/before/sfmn/9'
    save_dir = Path('find_template/results')

    # 递交任务
    img_truth_list = sorted(os.listdir(dir_truth))
    img_worse_list = sorted(os.listdir(dir_worse))
    img_better_list = sorted(os.listdir(dir_better))
    assert len(img_better_list) == len(img_truth_list) and len(
        img_better_list) == len(img_worse_list), '视频帧长度不相等'
    list_len = len(img_better_list)

    with ProcessPoolExecutor() as pool:
        for index in range(list_len):
            img_truth_path = osp.join(dir_truth, img_truth_list[index])
            img_worse_path = osp.join(dir_worse, img_worse_list[index])
            img_better_path = osp.join(dir_better, img_better_list[index])
            save_path = save_dir / osp.basename(dir_better) / (
                osp.basename(dir_better) + '-' +
                osp.splitext(img_truth_list[index])[0])

            path_config = {
                'img_truth_path': img_truth_path,
                'img_better_path': img_better_path,
                'img_worse_path': img_worse_path,
                'save_path': save_path
            }
            size_config = {'sub_width': sub_width, 'sub_height': sub_height}
            pool.submit(find_template, path_config, size_config, sliding_step,
                        mode)
        pool.shutdown(True)

    print("程序结束")
